package com.bzframework.ai.langchain4j;

import com.bzframework.ai.langchain4j.tool.SimpleFunctionTool;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.splitter.DocumentByParagraphSplitter;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.chat.StreamingChatModel;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import lombok.Getter;
import lombok.RequiredArgsConstructor;
import lombok.Setter;
import lombok.SneakyThrows;
import org.springframework.boot.autoconfigure.AutoConfiguration;
import org.springframework.context.annotation.Bean;
import org.springframework.core.io.ClassPathResource;

import java.util.List;

/**
 * QwenChatModelConfiguration qwen 配置类
 */
@Getter
@Setter
@AutoConfiguration
@RequiredArgsConstructor
public class QwenChatModelConfiguration {

    private final ChatModel qwenChatModel;

    private final StreamingChatModel qwenStreamingChatModel;

    private final EmbeddingModel qwenEmbeddingModel;

    @Bean
    public LangChain4jAiService langChain4jAiService(ContentRetriever contentRetriever) {
        MessageWindowChatMemory messageWindowChatMemory = MessageWindowChatMemory.withMaxMessages(10);
        return AiServices.builder(LangChain4jAiService.class)
                         .chatModel(qwenChatModel)
                         .streamingChatModel(qwenStreamingChatModel)
                         // 针对所有会话
                         .chatMemory(messageWindowChatMemory)
                         // 针对每一个用户实现会话记忆
                         .chatMemoryProvider(memoryId -> MessageWindowChatMemory.withMaxMessages(10))
                         // RAG 搜索
                         .contentRetriever(contentRetriever)
                         .tools(new SimpleFunctionTool())
                         .build();

    }

    @Bean
    @SneakyThrows
    public ContentRetriever contentRetriever(EmbeddingStore<TextSegment> embeddingStore) {
        // Rag相关配置 标准实现
        // 1.加载文档
        ClassPathResource resource = new ClassPathResource("langchain4j/docs");
        String filePath = resource.getFile().getAbsolutePath();
        List<Document> documents = FileSystemDocumentLoader.loadDocuments(filePath);
        // 2.文档切割，每个文档按照段落进行区分，最大1000，允许重叠200
        DocumentByParagraphSplitter documentByParagraphSplitter =
                new DocumentByParagraphSplitter(1000, 200);
        // 3.自定义文档加载器，把文档转换成向量向量并存储到向量数据库中
        EmbeddingStoreIngestor ingestor =
                EmbeddingStoreIngestor.builder()
                                      .documentSplitter(documentByParagraphSplitter)
                                      // 文档碎片添加文档名称作为原信息，提高文档质量
                                      .textSegmentTransformer(
                                              textSegment ->
                                                      TextSegment.from(textSegment.metadata().getString("file_name") + "\n" + textSegment.text(),
                                                                       textSegment.metadata())
                                      )
                                      .embeddingModel(qwenEmbeddingModel)
                                      .embeddingStore(embeddingStore)
                                      .build();
        ingestor.ingest(documents);
        return EmbeddingStoreContentRetriever.builder()
                                             // 将用户信息转换成向量类型的数据
                                             .embeddingModel(qwenEmbeddingModel)
                                             // 向量数据存储源
                                             .embeddingStore(embeddingStore)
                                             .maxResults(5) // 最大返回结果数
                                             .minScore(0.75) // 最小匹配分数
                                             .build();
    }

    @Bean
    public EmbeddingStore<TextSegment> embeddingStore() {
        return new InMemoryEmbeddingStore<>();
    }

}
